import logging
import socket
import time

logger = logging.getLogger(__name__)


class TCPConnectionError(Exception):
    pass


class TCPClient(object):
    """
    Basic networking client for TCP connections. Errors occurred during
    networking operations are raised as TCPConnectionError.

    Messages are sent and received as raw bytes without any length prefix.
    """

    def __init__(self, host: str, port: int, timeout: float = 10.0):
        self._host = host
        self._port = port
        self._timeout = timeout
        self._socket = None
        self._logprefix = "(%s:%s) " % (self._host, self._port)
        self._buffer_size = 10240  # Maximum receive buffer size

    def connect(self, connection_attempts: int = 10):
        """Try to establish a connection to the given host:port."""
        connection_attempts = max(1, connection_attempts)
        error = None
        for attempt in range(1, connection_attempts + 1):
            try:
                self._socket = socket.create_connection(
                    address=(self._host, self._port), timeout=self._timeout
                )
                self._socket.settimeout(self._timeout)
                logger.debug("%sconnected", self._logprefix)
                return
            except socket.error as exception:
                error = exception
                logger.debug(
                    "%sconnection attempt %d: %s", self._logprefix, attempt, error
                )
                time.sleep(1)
        self._reraise_exception_as_tcp_error("failed to connect", error)

    def disconnect(self):
        """Disconnect any active connection."""
        if self._socket is not None:
            logger.debug("%sdisconnecting", self._logprefix)
            self._socket.close()
            self._socket = None

    def connected(self):
        """Return whether there is an active connection."""
        return self._socket is not None

    def write(self, message: bytes):
        """Send message to the server."""
        if self._socket is None:
            raise TCPConnectionError(self._logprefix + "not connected")

        header = str(len(message)).zfill(4).encode("utf-8")
        try:
            self._socket.sendall(header + message)
        except socket.error as exception:
            self._reraise_exception_as_tcp_error("failed to write data", exception)

    def read(self):
        """Read data from the server."""
        header = self._read_n(4)
        if not header:
            raise TCPConnectionError(self._logprefix + "connection closed")
        length = int(header.decode("utf-8"))
        data = self._read_n(length)
        return data

    def _read_n(self, length):
        """Read n bytes from the socket"""
        if self._socket is None:
            raise TCPConnectionError(self._logprefix + "not connected")

        buf = bytes()
        while length > 0:
            try:
                data = self._socket.recv(length)
            except socket.error as exception:
                self._reraise_exception_as_tcp_error("failed to read data", exception)

            if not data:
                raise TCPConnectionError(self._logprefix + "connection closed")

            buf += data
            length -= len(data)

        return buf

    def _reraise_exception_as_tcp_error(self, message, exception):
        raise TCPConnectionError("%s%s: %s" % (self._logprefix, message, exception))
